package com.huan.pdf.utils;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * pdf字节流转换成图片字节流
 *
 * @author admin
 */
public final class Pdf2PicUtil {

    private static final int DPI = 150;
    private static final String IMG_TYPE = "jpg";

    /**
     * 单页 PDF → 图片字节
     */
    public static byte[] toImageBytes(byte[] pdfBytes) throws IOException {
        try (PDDocument doc = PDDocument.load(pdfBytes)) {
            BufferedImage bi = new PDFRenderer(doc).renderImageWithDPI(0, DPI);
            try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
                ImageIO.write(bi, IMG_TYPE, out);
                return out.toByteArray();
            }
        }
    }

    /**
     * 多页 PDF → 图片字节列表
     */
    public static List<byte[]> toImageBytesList(byte[] pdfBytes) throws IOException {
        List<byte[]> list = new ArrayList<>();
        try (PDDocument doc = PDDocument.load(pdfBytes)) {
            PDFRenderer renderer = new PDFRenderer(doc);
            for (int i = 0; i < doc.getNumberOfPages(); i++) {
                BufferedImage bi = renderer.renderImageWithDPI(i, DPI);
                try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
                    ImageIO.write(bi, IMG_TYPE, out);
                    list.add(out.toByteArray());
                }
            }
        }
        return list;
    }

    /**
     * 多页 PDF → zip 字节（内存完成，不落盘）
     *
     * @param pdfBytes 原始 PDF
     * @param fileName 不含扩展名的文件名（用于 zip 包内图片前缀）
     */
    public static byte[] toZipBytes(byte[] pdfBytes, String fileName) throws IOException {
        try (PDDocument doc = PDDocument.load(pdfBytes);
             ByteArrayOutputStream zipBuff = new ByteArrayOutputStream();
             ZipOutputStream zos = new ZipOutputStream(zipBuff)) {

            PDFRenderer renderer = new PDFRenderer(doc);
            int pages = doc.getNumberOfPages();

            for (int i = 0; i < pages; i++) {
                BufferedImage bi = renderer.renderImageWithDPI(i, DPI);
                // 单张图片转字节
                byte[] imgBytes;
                try (ByteArrayOutputStream imgOut = new ByteArrayOutputStream()) {
                    ImageIO.write(bi, IMG_TYPE, imgOut);
                    imgBytes = imgOut.toByteArray();
                }
                // 写入 zip
                String entryName = String.format("%s_%04d.%s", fileName, i + 1, IMG_TYPE);
                zos.putNextEntry(new ZipEntry(entryName));
                zos.write(imgBytes);
                zos.closeEntry();
            }
            zos.finish();
            return zipBuff.toByteArray();
        }
    }
}